import os
import time
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.font_manager import FontProperties
from tqdm import tqdm, trange
from colorama import Fore, Back, Style, init
from tabulate import tabulate
from datetime import datetime
import json
import psutil
import torch
import gc
import GPUtil

# 初始化 colorama
init(autoreset=True)

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

class PrettyTrainer:
    """美观的模型训练输出管理器"""
    
    def __init__(self, model_name, epochs, save_dir="./runs", plot_interval=5):
        """初始化训练管理器
        
        Args:
            model_name (str): 模型名称
            epochs (int): 训练轮数
            save_dir (str): 保存路径
            plot_interval (int): 图表绘制间隔轮数
        """
        self.model_name = model_name
        self.epochs = epochs
        self.start_time = time.time()
        self.plot_interval = plot_interval
        self.metrics_history = {
            'train_loss': [],
            'val_loss': [],
            'train_acc': [],
            'val_acc': [],
            'learning_rate': [],
            'epoch_time': []
        }
        self.best_val_acc = 0
        self.best_epoch = 0
        self.early_stop_counter = 0
        
        # 创建保存目录
        self.timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.run_name = f"{model_name}_{self.timestamp}"
        self.save_dir = os.path.join(save_dir, self.run_name)
        os.makedirs(self.save_dir, exist_ok=True)
        
        # 显示初始化信息
        self._print_header()
        
        # 记录初始系统信息
        self._log_system_info()
    
    def _log_system_info(self):
        """记录系统信息"""
        system_info = {
            "CPU使用率": f"{psutil.cpu_percent()}%",
            "内存使用率": f"{psutil.virtual_memory().percent}%",
            "可用内存": f"{psutil.virtual_memory().available / (1024**3):.2f} GB",
            "总内存": f"{psutil.virtual_memory().total / (1024**3):.2f} GB"
        }
        
        if torch.cuda.is_available():
            system_info["GPU设备"] = torch.cuda.get_device_name(0)
            system_info["GPU显存总量"] = f"{torch.cuda.get_device_properties(0).total_memory / (1024**3):.2f} GB"
            system_info["GPU显存使用"] = f"{torch.cuda.memory_allocated(0) / (1024**3):.2f} GB"
            system_info["GPU显存占用率"] = f"{torch.cuda.memory_allocated(0) / torch.cuda.get_device_properties(0).total_memory * 100:.2f}%"
        
        print(f"\n{Fore.CYAN}{Style.BRIGHT}🖥️ 系统资源信息 🖥️{Style.RESET_ALL}\n")
        system_table = [[k, v] for k, v in system_info.items()]
        print(tabulate(system_table, headers=["指标", "数值"], tablefmt="fancy_grid"))
    
    def _print_header(self):
        """打印训练开始的标题"""
        print("\n" + "="*80)
        print(f"{Fore.CYAN}{Style.BRIGHT}✨ 深度学习训练任务 ✨{Style.RESET_ALL}".center(90))
        print(f"{Fore.YELLOW}模型: {Style.BRIGHT}{self.model_name}{Style.RESET_ALL}".center(80))
        print(f"{Fore.YELLOW}保存路径: {Style.BRIGHT}{self.save_dir}{Style.RESET_ALL}".center(90))
        print("="*80 + "\n")
    
    def process_data(self, dataset_name, total_steps, additional_info=None):
        """显示数据处理进度
        
        Args:
            dataset_name (str): 数据集名称
            total_steps (int): 总步数
            additional_info (dict, optional): 附加信息
        """
        print(f"\n{Fore.GREEN}>>> {Style.BRIGHT}数据处理: {dataset_name}{Style.RESET_ALL}")
        
        if additional_info:
            for key, value in additional_info.items():
                print(f"{Fore.WHITE}• {key}: {Fore.CYAN}{value}{Style.RESET_ALL}")
        
        # 创建一个带有描述的进度条
        pbar = tqdm(total=total_steps, 
                   desc=f"{Fore.WHITE}处理进度", 
                   bar_format="{l_bar}{bar:30}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]",
                   ncols=100)
        
        return pbar
    
    def update_data_progress(self, pbar, step, message=None):
        """更新数据处理进度
        
        Args:
            pbar (tqdm): 进度条对象
            step (int): 步数
            message (str, optional): 显示的消息
        """
        if message:
            pbar.set_postfix_str(message)
        pbar.update(step)
    
    def data_summary(self, data_info):
        """展示数据处理摘要
        
        Args:
            data_info (dict): 数据信息字典，包含各种数据统计
        """
        print(f"\n{Fore.BLUE}{Style.BRIGHT}📊 数据摘要 📊{Style.RESET_ALL}\n")
        
        # 转换为表格格式
        table_data = []
        for key, value in data_info.items():
            table_data.append([key, str(value)])
        
        # 使用 tabulate 打印美观的表格
        print(tabulate(table_data, headers=["指标", "数值"], tablefmt="fancy_grid"))
        print("\n")
        
        # 如果有类别分布，绘制图表
        if 'class_distribution' in data_info and isinstance(data_info['class_distribution'], dict):
            plt.figure(figsize=(10, 6))
            classes = list(data_info['class_distribution'].keys())
            counts = list(data_info['class_distribution'].values())
            plt.bar(classes, counts, color='skyblue')
            plt.title('类别分布')
            plt.xlabel('类别')
            plt.ylabel('样本数')
            plt.xticks(rotation=45)
            plt.tight_layout()
            plt.savefig(os.path.join(self.save_dir, "class_distribution.png"))
            plt.close()
            print(f"{Fore.GREEN}✓ 类别分布图已保存到: {self.save_dir}/class_distribution.png{Style.RESET_ALL}")
    
    def model_summary(self, model_info):
        """展示模型结构摘要
        
        Args:
            model_info (dict): 模型信息字典
        """
        print(f"\n{Fore.MAGENTA}{Style.BRIGHT}🏗️ 模型架构摘要 🏗️{Style.RESET_ALL}\n")
        
        # 计算参数量
        total_params = model_info.get('total_params', 0)
        trainable_params = model_info.get('trainable_params', 0)
        
        params_info = [
            ["模型类型", model_info.get('model_type', '-')],
            ["总参数量", f"{total_params:,}"],
            ["可训练参数", f"{trainable_params:,}"],
            ["非训练参数", f"{total_params - trainable_params:,}"],
            ["优化器", model_info.get('optimizer', '-')],
            ["学习率", model_info.get('learning_rate', '-')],
            ["损失函数", model_info.get('loss_function', '-')]
        ]
        
        print(tabulate(params_info, headers=["属性", "值"], tablefmt="fancy_grid"))
        
        # 如果有详细的层信息，打印层结构
        if 'layers' in model_info and isinstance(model_info['layers'], list):
            print(f"\n{Fore.YELLOW}层结构:{Style.RESET_ALL}")
            
            # 过滤掉不重要的层
            important_layers = [
                layer for layer in model_info['layers']
                if any(key in layer['name'].lower() for key in 
                      ['conv', 'linear', 'pool', 'norm', 'dropout', 'relu', 'silu', 'fc'])
            ]
            
            # 创建层信息表格
            layers_table = []
            for i, layer in enumerate(important_layers):
                name = layer['name']
                layer_type = layer['type']
                shape = layer.get('output_shape', '-')
                params = layer.get('params', 0)
                
                # 为不同类型的层使用不同的颜色
                if 'conv' in layer_type.lower():
                    name = f"{Fore.CYAN}{name}{Style.RESET_ALL}"
                elif 'linear' in layer_type.lower() or 'fc' in name.lower():
                    name = f"{Fore.GREEN}{name}{Style.RESET_ALL}"
                elif 'norm' in layer_type.lower():
                    name = f"{Fore.YELLOW}{name}{Style.RESET_ALL}"
                elif 'pool' in layer_type.lower():
                    name = f"{Fore.MAGENTA}{name}{Style.RESET_ALL}"
                
                layers_table.append([
                    i + 1,
                    name,
                    layer_type,
                    shape,
                    f"{params:,}" if params > 0 else "-"
                ])
            
            print(tabulate(
                layers_table,
                headers=["层", "名称", "类型", "输出形状", "参数量"],
                tablefmt="pretty",
                colalign=("right", "left", "left", "right", "right")
            ))
    
    def start_epoch(self, epoch, learning_rate=None):
        """开始新的训练轮次
        
        Args:
            epoch (int): 当前轮次
            learning_rate (float, optional): 当前学习率
        
        Returns:
            dict: 本轮的进度条和时间记录
        """
        epoch_start_time = time.time()
        progress = {}
        
        # 添加学习率到历史记录
        if learning_rate is not None:
            self.metrics_history['learning_rate'].append(learning_rate)
        
        print(f"\n{Fore.MAGENTA}{Style.BRIGHT}🔄 第 {epoch+1}/{self.epochs} 轮训练开始 🔄{Style.RESET_ALL}")
        if learning_rate is not None:
            print(f"{Fore.CYAN}当前学习率: {learning_rate}{Style.RESET_ALL}")
        
        # 记录内存使用情况
        if torch.cuda.is_available():
            mem_allocated = torch.cuda.memory_allocated(0) / (1024**3)
            mem_reserved = torch.cuda.memory_reserved(0) / (1024**3)
            print(f"{Fore.WHITE}GPU内存: 已分配 {mem_allocated:.2f} GB, 已保留 {mem_reserved:.2f} GB{Style.RESET_ALL}")
        
        # 创建单个进度条字典，而不是多个进度条
        progress['train'] = None
        progress['val'] = None
        
        return {'progress': progress, 'start_time': epoch_start_time}
    
    def update_batch_metrics(self, progress_bar, batch_metrics, batch_idx, total_batches, phase='train'):
        """更新单个批次的训练指标
        
        Args:
            progress_bar (tqdm): 进度条对象
            batch_metrics (dict): 批次指标
            batch_idx (int): 批次索引
            total_batches (int): 总批次数
            phase (str): 当前阶段 ('train' 或 'val')
        """
        # 格式化批次指标
        metrics_str = f"loss: {batch_metrics['loss']:.4f}, acc: {batch_metrics['acc']:.4f}"
        
        # 如果进度条为None，创建新的进度条
        if progress_bar is None:
            desc = f"{Fore.GREEN}训练{Style.RESET_ALL}" if phase == 'train' else f"{Fore.BLUE}验证{Style.RESET_ALL}"
            progress_bar = tqdm(total=total_batches, 
                              desc=desc, 
                              bar_format="{desc}: {percentage:3.0f}%|{bar:30}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}] {postfix}",
                              ncols=120,
                              position=0,
                              leave=True,
                              dynamic_ncols=True)  # 允许动态调整宽度
        
        # 更新进度条
        progress_bar.n = batch_idx
        progress_bar.refresh()
        
        # 更新显示的指标
        progress_bar.set_postfix_str(metrics_str)
        
        # 在最后一个batch时确保显示100%并关闭
        if batch_idx >= total_batches - 1:
            progress_bar.n = total_batches
            progress_bar.refresh()
            progress_bar.close()
            
        return progress_bar
    
    def update_metrics(self, epoch, metrics, save_plot=False):
        """更新并显示训练指标
        
        Args:
            epoch (int): 当前轮次
            metrics (dict): 包含train_loss, val_loss, train_acc, val_acc等指标的字典
            save_plot (bool): 是否保存图表
        """
        # 保存指标历史
        for key, value in metrics.items():
            if key in self.metrics_history:
                self.metrics_history[key].append(value)
        
        # 记录最佳模型
        if metrics.get('val_acc', 0) > self.best_val_acc:
            self.best_val_acc = metrics.get('val_acc', 0)
            self.best_epoch = epoch
            improved = f"{Fore.GREEN}✓ 最佳模型!{Style.RESET_ALL}"
            self.early_stop_counter = 0
        else:
            improved = ""
            self.early_stop_counter += 1
        
        # 创建一个漂亮的表格来显示当前轮次的指标
        epoch_data = [
            ["训练损失", f"{Fore.CYAN}{metrics.get('train_loss', 0):.4f}{Style.RESET_ALL}"],
            ["验证损失", f"{Fore.YELLOW}{metrics.get('val_loss', 0):.4f}{Style.RESET_ALL}"],
            ["训练准确率", f"{Fore.CYAN}{metrics.get('train_acc', 0)*100:.2f}%{Style.RESET_ALL}"],
            ["验证准确率", f"{Fore.YELLOW}{metrics.get('val_acc', 0)*100:.2f}%{Style.RESET_ALL}"]
        ]
        
        # 添加其他指标
        for key, value in metrics.items():
            if key not in ['train_loss', 'val_loss', 'train_acc', 'val_acc']:
                if isinstance(value, float):
                    epoch_data.append([key, f"{value:.4f}"])
                else:
                    epoch_data.append([key, str(value)])
        
        print(f"\n{Fore.GREEN}{Style.BRIGHT}📈 第 {epoch+1} 轮指标 📈{Style.RESET_ALL} {improved}")
        print(tabulate(epoch_data, headers=["指标", "数值"], tablefmt="fancy_grid"))
        
        # 打印最佳模型信息
        print(f"{Fore.YELLOW}最佳验证准确率: {self.best_val_acc*100:.2f}% (轮次 {self.best_epoch+1}){Style.RESET_ALL}")
        
        # 如果需要保存图表
        if save_plot:
            self._plot_training_history()
    
    def _plot_training_history(self):
        """绘制训练历史图表"""
        plt.figure(figsize=(15, 5))
        
        # 损失曲线
        plt.subplot(1, 2, 1)
        plt.plot([float(x) for x in self.metrics_history['train_loss']], label='训练损失', color='#2ecc71')
        plt.plot([float(x) for x in self.metrics_history['val_loss']], label='验证损失', color='#e74c3c')
        plt.title('损失曲线')
        plt.xlabel('轮次')
        plt.ylabel('损失')
        plt.legend()
        plt.grid(True, linestyle='--', alpha=0.7)
        
        # 准确率曲线
        plt.subplot(1, 2, 2)
        plt.plot([float(x)*100 for x in self.metrics_history['train_acc']], label='训练准确率', color='#2ecc71')
        plt.plot([float(x)*100 for x in self.metrics_history['val_acc']], label='验证准确率', color='#e74c3c')
        plt.title('准确率曲线')
        plt.xlabel('轮次')
        plt.ylabel('准确率 (%)')
        plt.legend()
        plt.grid(True, linestyle='--', alpha=0.7)
        
        plt.tight_layout()
        save_path = os.path.join(self.save_dir, f'training_history_{len(self.metrics_history["train_loss"])}.png')
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.close()
        
        print(f"\n{Fore.GREEN}✓ 训练历史图表已保存到: {save_path}{Style.RESET_ALL}")